from torch.utils.data import DataLoader, Dataset
import numpy as np
import random
from PIL import Image
from PIL import ImageFile
ImageFile.LOAD_TRUNCATED_IMAGES = True
import PIL.ImageOps
import pandas as pd
import os

pic = pd.read_csv(r'/media/wz209/a29353b7-1090-433f-b452-b4ce827adb17/sugurs/Dataset/release_v0/meta/meta.csv')
pic.columns = ["Col1", "Col2", "Col3", "Col4", "Col5", "Col6", "Col7",
               "Col8", "Col9", "Col10", "Col11", "Col12",
               "Col13", "Col14", "Col15", "Col16", "Col17","Col18","Col19"]
Y = pic[["Col1", "Col2", "Col4", "Col5", "Col6", "Col7", "Col8", "Col9", "Col10", "Col16", "Col17"]]
pic_path = r"/media/wz209/a29353b7-1090-433f-b452-b4ce827adb17/sugurs/Dataset/release_v0/images"

# pic = pd.read_csv(
#     r'/media/citaa/e8d7592e-674f-496d-a4f2-14ce0f3d8098/dataset/data/handle.csv'
# )
# 19
# pic.columns = ["Col1", "Col2", "Col3", "Col4", "Col5", "Col6", "Col7", "Col8", "Col9", "Col10", "Col11", "Col12",
#                "Col13", "Col14", "Col15", "Col16"]
# Y = pic[["Col1", "Col16"]]
# pic_path = r"/media/citaa/e8d7592e-674f-496d-a4f2-14ce0f3d8098/dataset/data/data"
# dict = {}
# for i in range(0,len(Y)-1):
#     dict[Y[i][0]] = i

Y = np.array(Y)

class Dataset(Dataset):
    def __init__(self, imagepath, transform=None, should_invert=True):
        self.imageFolderDataset = imagepath

        self.imageFolderDataset = pd.read_csv(imagepath)
        self.imageFolderDataset.columns = ["Col1"]
        self.data = self.imageFolderDataset[["Col1"]]
        self.data = np.array(self.data)

        self.transform = transform
        self.should_invert = should_invert

    def __getitem__(self, index):

        num = random.randrange(0, len(self.data))
        key = self.data[num][0] - 1
        # num行对应到num+2  1->3 self.data[0][0] = 16
        img1_tuple = Y[key][9]
        # 对应第key+1个case
        img2_tuple = Y[key][10]

        label = []
        name = Y[key][2] # PN 0
        if name == 'absent':
            label.append(0)
        elif name == 'typical':
            label.append(1)
        else:
            label.append(2)

        name = Y[key][3] # STR 1
        if name == 'absent':
            label.append(0)
        elif name == 'regular':
            label.append(1)
        else:
            label.append(2)

        name = Y[key][4] # PIG 2
        if name == 'absent':
            label.append(0)
        elif name in ['diffuse regular', 'localized regular']:
            label.append(1)
        else:
            label.append(2)

        name = Y[key][5] # RS 3
        if name == 'absent':
            label.append(0)
        else:
            label.append(1)

        name = Y[key][6] # DaG 4
        if name == 'absent':
            label.append(0)
        elif name == 'regular':
            label.append(1)
        else:
            label.append(2)

        name = Y[key][7] # BWV 5
        if name == 'absent':
            label.append(0)
        else:
            label.append(1)

        name = Y[key][8] # VS 6
        if name == 'absent':
            label.append(0)
        elif name in ['arborizing', 'comma', 'hairpin', 'within regression', 'wreath']:
            label.append(1)
        else:
            label.append(2)

        name = Y[key][1] # DIAG 7
        list1 = ['basal cell carcinoma']
        list2 = ['blue nevus', 'clark nevus', 'combined nevus', 'congenital nevus', 'dermal nevus', 'recurrent nevus',
                 'reed or spitz nevus']
        list3 = ['melanoma (in situ)', 'melanoma', 'melanoma (less than 0.76 mm)', 'melanoma (0.76 to 1.5 mm)',
                 'melanoma (more than 1.5 mm)', 'melanoma metastasis']
        list4 = ['dermatofibroma', 'lentigo', 'melanosis', 'miscellaneous', 'vascular lesion']
        list5 = ['seborrheic keratosis']
        if name in list1:
            label.append(0)
        elif name in list2:
            label.append(1)
        elif name in list3:
            label.append(2)
        elif name in list4:
            label.append(3)
        else:
            label.append(4)

        train1 = img1_tuple.split('/')
        train1_ = os.path.join(train1[0], train1[1])
        img1_tuple = os.path.join(pic_path, train1_)

        train2 = img2_tuple.split('/')
        train2_ = os.path.join(train2[0], train2[1])
        img2_tuple = os.path.join(pic_path, train2_)

        img0 = Image.open(img1_tuple)
        img1 = Image.open(img2_tuple)

        if self.should_invert:
            img0 = PIL.ImageOps.invert(img0)
            img1 = PIL.ImageOps.invert(img1)

        if self.transform is not None:
            img0 = self.transform(img0)
            img1 = self.transform(img1)

        return img0, img1, label

    def __len__(self):
        return len(self.data)
'''
class Dataset(Dataset):
    def __init__(self, imagepath, transform=None, should_invert=True):
        self.imageFolderDataset = imagepath

        self.imageFolderDataset = pd.read_csv(imagepath)
        self.imageFolderDataset.columns = ["Col1"]
        self.data = self.imageFolderDataset[["Col1"]]
        self.data = np.array(self.data)

        self.transform = transform
        self.should_invert = should_invert

    def __getitem__(self, index):

        num = random.randrange(0, len(self.data)-1)
        key = self.data[num][0]
        #print(num,key,dict[key],Y[dict[key]][1])

        label = int(Y[dict[key]][1]) - 1
        if label == 5:
            label = 3

        train1 = str(key)
        train1 = os.path.join("cli",train1)
        train1 = os.path.join(str(int(key)), train1)
        img1_tuple = os.path.join(pic_path, train1)
        if os.path.exists(img1_tuple + "_1.JPG"):
            img1_tuple = img1_tuple + "_1.JPG"
        elif os.path.exists(img1_tuple + "_1.jpg"):
            img1_tuple = img1_tuple + "_1.jpg"
        elif os.path.exists(img1_tuple + "_1.jpeg"):
            img1_tuple = img1_tuple + "_1.jpeg"
        elif os.path.exists(img1_tuple + "_1.PNG"):
            img1_tuple = img1_tuple + "_1.PNG"

        train2 = str(key)
        train2 = os.path.join("der", train2)
        train2 = os.path.join(str(int(key)), train2)
        img2_tuple = os.path.join(pic_path, train2)
        if os.path.exists(img2_tuple + "_1.JPG"):
            img2_tuple = img2_tuple + "_1.JPG"
        elif os.path.exists(img2_tuple + "_1.jpg"):
            img2_tuple = img2_tuple + "_1.jpg"
        elif os.path.exists(img2_tuple + "_1.jpeg"):
            img2_tuple = img2_tuple + "_1.jpeg"
        elif os.path.exists(img2_tuple + "_1.PNG"):
            img2_tuple = img2_tuple + "_1.PNG"

        img0 = Image.open(img1_tuple)
        img1 = Image.open(img2_tuple)

        if self.should_invert:
            img0 = PIL.ImageOps.invert(img0)
            img1 = PIL.ImageOps.invert(img1)

        if self.transform is not None:
            img0 = self.transform(img0)
            img1 = self.transform(img1)

        return img0, img1, label

    def __len__(self):
        return len(self.data)
'''